#!/usr/bin/python
# diff for mongo collection structure between mongo databases.
# author tom
# 2021年3月8日
import json
import operator
import os
import sys
import traceback
import time
from configparser import ConfigParser
from datetime import datetime

from pymongo import MongoClient
from pymongo.collection import Collection
from tqdm import tqdm


def p(msg):
    print(f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")} => {msg}')


def __fill_absent_config(_collection: list):
    absent = 4 - len(_collection)
    if absent > 0:
        for i in range(len(_collection) - 1, 3):
            _collection.append('')
    return _collection


def __parse_diff_collections(collections_str):
    _collections = []
    collections_str = collections_str.strip()
    collection_arr = collections_str.split(',')
    for c_str in collection_arr:
        _cfg = __fill_absent_config(c_str.split(":"))
        _collection = {
            'name': _cfg[0],
            'sort': ((_cfg[1], int(_cfg[2]) if len(_cfg[2]) > 0 else -1)
                     if len(_cfg[1]) > 0 else None),
            'limit': int(_cfg[-1]) if len(_cfg[-1]) > 0 else None
        }
        _collections.append(_collection)
    return _collections


def __close(closable):
    if closable:
        closable.close()


def __diff(s1, s2, _diff, c1_info, c2_info):
    merge_s = s1.copy()
    merge_s.update(s2)
    for k, v in merge_s.items():
        if s1.__contains__(k) and not s2.__contains__(k):
            if 'type' in s1[k] and ('dict' in s1[k]['type'] or 'list' in s1[k]['type']) and len(s1[k]['children']) > 0:
                _diff[k] = {'type': s1[k]['type'], 'children': {}, 'owner': c1_info}
                __diff(s1[k]['children'], {}, _diff[k]['children'], c1_info, c2_info)
            else:
                _diff[k] = s1[k]
                _diff[k]['owner'] = c1_info
        elif s2.__contains__(k) and not s1.__contains__(k):
            if 'type' in s2[k] and ('dict' in s2[k]['type'] or 'list' in s2[k]['type']) and len(s2[k]['children']) > 0:
                _diff[k] = {'type': s2[k]['type'], 'children': {}, 'owner': c2_info}
                __diff({}, s2[k]['children'], _diff[k]['children'], c1_info, c2_info)
            else:
                _diff[k] = s2[k]
                _diff[k]['owner'] = c2_info
        elif k in s1 and k in s2 and operator.ne(type(s1[k]), type(s2[k])):
            _diff[k] = {}
            if 'type' not in _diff[k]:
                _diff[k]['type'] = set()
            _diff[k]['type'].add('{}[{}]'.format(c1_info, str(type(s1[k])).split("'")[-2]))
            _diff[k]['type'].add('{}[{}]'.format(c2_info, str(type(s2[k])).split("'")[-2]))
            _diff[k]['owner'] = c1_info + ',' + c2_info
        elif issubclass(type(s1[k]), dict):
            _diff[k] = {}
            __diff(s1[k], s2[k], _diff[k], c1_info, c2_info)
            if len(_diff[k]) <= 0:
                del _diff[k]


def __list_parse(_list: list, structure: dict):
    for _item in _list:
        if isinstance(_item, list):
            array_index_name = '@array-' + str(len(set(filter(lambda x: x.startswith('@array'), structure.keys()))))
            structure[array_index_name] = {
                'type': {'list', },
                'children': {}
            }
            __list_parse(_item, structure[array_index_name]['children'])
        elif isinstance(_item, dict):
            __dict_parse(_item, structure)


def __dict_parse(item: dict, structure: dict):
    for col_name, col_value in item.items():
        # p('col_name: %s, col_value: %s, value_type: %s' % (col_name, col_value, type(col_value)))
        # <Class 'str'> => str
        col_type = str(type(col_value)).split("'")[-2]
        if col_name not in structure:
            structure[col_name] = {
                'type': {col_type, },
                'children': {}
            }
        else:
            structure[col_name]['type'].add(col_type)
        if isinstance(col_value, list):
            __list_parse(col_value, structure[col_name]['children'])
        elif isinstance(col_value, dict):
            __dict_parse(col_value, structure[col_name]['children'])
        # if 'children' in structure[col_name] and len(structure[col_name]['children']) <= 0:
        #     del structure[col_name]['children']


def __get_db_info(_c: Collection, c_info):
    if _c:
        return '{}/{}/{}'.format(':'.join('{}'.format(v) for v in _c.database.client.address),
                                 _c.database.name, c_info['name'])
    return None


def __get_structure(_c: Collection, c_info):
    p(f'Analysing collection: {c_info["name"]}')
    _struct = {}
    # more fast than count_document()
    _count = _c.estimated_document_count()
    total = c_info['limit'] if c_info['limit'] and c_info['limit'] <= _count else _count
    # for item in tqdm(
    #         asyncio.as_completed(_c.find(limit=(c_info['limit'] if c_info['limit'] else 0),
    #                                                sort=([c_info['sort']] if c_info['sort'] is not None else None))),
    #         total=total,
    #         desc='[{}/{}]Analysing -> {}'.format(_c.database.client.address, _c.database.name, c_info['name'])
    # ):
    db_info = __get_db_info(_c, c_info)
    p(f'Database: {db_info}, collection: {c_info["name"]}, total: {_count}, limit: {total}')
    for item in tqdm(
            _c.find(limit=(c_info['limit'] if c_info['limit'] else 0),
                    sort=([c_info['sort']] if c_info['sort'] is not None else None)),
            total=total,
            desc=f"Analysing {db_info}"):
        __dict_parse(item, _struct)
    return db_info, _struct


def __diff_collection(c1, c2, cfg):
    struct_arr = {}
    c1_info = __get_db_info(c1, cfg)
    c2_info = __get_db_info(c2, cfg)
    struct_arr[c1_info] = {}
    struct_arr[c2_info] = {}
    if c1:
        db_info, struct_info = __get_structure(c1, cfg)
        struct_arr[db_info] = struct_info
    if c2:
        db_info, struct_info = __get_structure(c2, cfg)
        struct_arr[db_info] = struct_info
    diff_res = {}
    __diff(struct_arr[c1_info], struct_arr[c2_info], diff_res, c1_info, c2_info)
    return cfg['name'], diff_res


def diff_structure(config_parser: ConfigParser):
    diff_collections = __parse_diff_collections(config_parser.get('mongo', 'collections'))
    db_name_1 = config_parser.get('mongo', 'db.name.1') if config_parser.has_option('mongo', 'db.name.1') else None
    db_name_2 = config_parser.get('mongo', 'db.name.2') if config_parser.has_option('mongo', 'db.name.2') else None
    uri_1 = config_parser.get('mongo', 'db.uri.1') if config_parser.has_option('mongo', 'db.uri.1') else None
    uri_2 = config_parser.get('mongo', 'db.uri.2') if config_parser.has_option('mongo', 'db.uri.2') else None
    p(f'Diff collections: {diff_collections}')
    p(f'Diff databases: {uri_1},{db_name_1}\t{uri_2},{db_name_2}')
    client_1 = MongoClient(uri_1) if uri_1 else None
    client_2 = MongoClient(uri_2) if uri_2 else None
    try:
        diff_info = []
        for _cfg in tqdm(diff_collections, desc="Collection Structure Analysing"):
            t_collection, t_res = __diff_collection(client_1[db_name_1][_cfg['name']] if client_1 else None,
                                                    client_2[db_name_2][_cfg['name']] if client_2 else None,
                                                    _cfg)
            if t_res:
                diff_info.append({
                    'collection': t_collection,
                    'fields': t_res
                })
        return diff_info
    except RuntimeError:
        traceback.print_exc()
    finally:
        __close(client_1)
        __close(client_2)
    return None


def __print_structure_tree(structure, levels, output):
    for node in sorted(structure.keys()):
        if 'type' in structure[node]:
            output('|\t' * levels + '|____', node, '[' + structure[node]['owner'] + ']',
                   ':', ', '.join(sorted(list(structure[node]['type']))))
            # if levels > 0:
            #     print('\t' * levels + '|____', node, '[' + structure[node]['owner'] + ']',
            #           ':', ', '.join(sorted(list(structure[node]['type']))))
            # else:
            #     print('|\t' * levels + '|____', node, '[' + structure[node]['owner'] + ']',
            #           ':', ', '.join(sorted(list(structure[node]['type']))))
        else:
            output('|\t' * levels + '|____', node)
        if len(structure[node]['children']) == 0:
            continue
        __print_structure_tree(structure[node]['children'], levels + 1, output)


class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        # set not serializable, convert to list
        if isinstance(obj, set):
            return list(obj)
        else:
            return super().default(obj)


def __print_to_stream(diff_structs, output):
    for s in diff_structs:
        output('=' * 132)
        # index:[align][width]  align: ^center, <left, >right
        output('=' * 50, '{0:^30}'.format(s['collection']), '=' * 50)
        output('=' * 132)
        __print_structure_tree(s['fields'], 0, output)


def __print(diff_structs, to_file=False):
    old_stream = sys.stdout
    if to_file:
        sys.stdout = open('struct{}.txt'.format(datetime.now().strftime("%Y%m%d%H%M%S")), mode='w+', encoding='utf-8')
    __print_to_stream(diff_structs, print)
    if to_file:
        f = sys.stdout
        sys.stdout = old_stream
        __close(f)


if __name__ == '__main__':
    config_file = os.path.join(os.path.dirname(__file__), 'config.ini')
    config = ConfigParser()
    config.read(config_file, 'utf-8')
    p(f'Read config file: {config_file}')
    start_time = time.mktime(datetime.now().timetuple())
    diff_result = diff_structure(config)
    p(json.dumps(diff_result, cls=MyEncoder))
    __print(diff_result, config.getboolean('mongo', 'output.file', fallback=False))
    p(f'Diff time cost: {time.mktime(datetime.now().timetuple()) - start_time}s')
